import torch
import torch.nn as nn


class Linear(nn.Module):
    def __init__(self, in_feats, n_class, m_drop=False):
        super().__init__()
        self.classifier = torch.nn.Linear(in_feats, n_class)
        self.dropout = nn.Dropout(0.2)
        self.m_drop = m_drop

    def forward(self, inputs):
        if self.training and self.m_drop:
            logits = torch.stack([self.classifier(self.dropout(inputs)) for _ in range(5)]).mean(0)
        else:
            logits = self.classifier(self.dropout(inputs))
        return logits
